/*
 * Copyright (c) 2023, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

package jdk.httpclient.test.lib.common;

import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.util.Date;
import java.util.Objects;

import javax.net.ssl.KeyManager;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;

import sun.security.x509.CertificateExtensions;
import sun.security.x509.CertificateSerialNumber;
import sun.security.x509.CertificateValidity;
import sun.security.x509.CertificateVersion;
import sun.security.x509.CertificateX509Key;
import sun.security.x509.DNSName;
import sun.security.x509.GeneralName;
import sun.security.x509.GeneralNames;
import sun.security.x509.SubjectAlternativeNameExtension;
import sun.security.x509.X500Name;
import sun.security.x509.X509CertImpl;
import sun.security.x509.X509CertInfo;

/**
 * A utility for generating dynamic {@link java.security.KeyStore}s in tests. The keystores
 * generated by this utility are dynamic in the sense that they are generated as necessary and
 * don't require keys/certificates to be present on the filesystem. The generated keystores too
 * aren't saved to the filesystem, by this utility.
 */
public class DynamicKeyStoreUtil {

    private static final String PKCS12_KEYSTORE_TYPE = "PKCS12";

    /**
     * The default alias that will be used in the KeyStore generated by
     * {@link #generateKeyStore(String, String...)}
     */
    public static final String DEFAULT_ALIAS = "foobar-key-alias";

    private static final String DEFAULT_SUBJECT_OU = "foobar-org-unit";
    private static final String DEFAULT_SUBJECT_ON = "foobar-org-name";
    private static final String DEFAULT_SUBJECT_COUNTRY = "US";


    /**
     * Generates a PKCS12 type {@link KeyStore} which has one
     * {@link KeyStore#getKey(String, char[]) key entry} which corresponds to a newly generated
     * {@link java.security.PrivateKey} and accompanied by a newly generated self-signed certificate,
     * certifying the corresponding public key.
     * <p>
     * The newly generated {@link X509Certificate} certificate will use the {@code certSubject}
     * as the certificate's subject. If the {@code certSubjectAltNames} is non-null then
     * the certificate will be created with a
     * {@link sun.security.x509.SubjectAlternativeNameExtension subject alternative name extension}
     * which will include the {@code certSubject} and each of the {@code certSubjectAltNames} as
     * subject alternative names (represented as DNS names)
     * <p>
     * The generated KeyStore won't be password protected
     *
     * @param certSubject         The subject to be used for the newly generated certificate
     * @param certSubjectAltNames Optional subject alternative names to be used in the generated
     *                            certificate
     * @return The newly generated KeyStore
     * @throws NullPointerException If {@code certSubject} is null
     */
    public static KeyStore generateKeyStore(final String certSubject, final String... certSubjectAltNames)
            throws Exception {
        Objects.requireNonNull(certSubject);
        final SecureRandom secureRandom = new SecureRandom();
        final KeyPair keyPair = generateRSAKeyPair(secureRandom);
        final X509Certificate cert = generateCert(keyPair, secureRandom, certSubject, certSubjectAltNames);
        final KeyStore.Builder keystoreBuilder = KeyStore.Builder.newInstance(PKCS12_KEYSTORE_TYPE,
                null, new KeyStore.PasswordProtection(null));
        final KeyStore keyStore = keystoreBuilder.getKeyStore();
        // write a private key (with cert chain for the public key)
        final char[] keyPassword = null;
        keyStore.setKeyEntry(DEFAULT_ALIAS, keyPair.getPrivate(), keyPassword, new Certificate[]{cert});
        return keyStore;
    }

    /**
     * Generates a PKCS12 type {@link KeyStore} which has one
     * {@link KeyStore#getKey(String, char[]) key entry} which corresponds to the {@code privateKey}
     * and accompanied by the {@code certChain} certifying the corresponding public key.
     * <p>
     * The generated KeyStore won't be password protected
     *
     * @param privateKey The private key to include in the keystore
     * @param certChain  The certificate chain
     * @return The newly generated KeyStore
     * @throws NullPointerException If {@code privateKey} or {@code certChain} is null
     */
    public static KeyStore generateKeyStore(final PrivateKey privateKey, final Certificate[] certChain)
            throws Exception {
        Objects.requireNonNull(privateKey);
        Objects.requireNonNull(certChain);
        final KeyStore.Builder keystoreBuilder = KeyStore.Builder.newInstance(PKCS12_KEYSTORE_TYPE,
                null, new KeyStore.PasswordProtection(null));
        final KeyStore keyStore = keystoreBuilder.getKeyStore();
        // write a private key (with cert chain for the public key)
        final char[] keyPassword = null;
        keyStore.setKeyEntry(DEFAULT_ALIAS, privateKey, keyPassword, certChain);
        return keyStore;
    }

    /**
     * Creates and returns a new PKCS12 type keystore without any entries in the keystore
     *
     * @return The newly created keystore
     * @throws Exception if any exception occurs during keystore generation
     */
    public static KeyStore generateBlankKeyStore() throws Exception {
        final KeyStore.Builder keystoreBuilder = KeyStore.Builder.newInstance(PKCS12_KEYSTORE_TYPE,
                null, new KeyStore.PasswordProtection(null));
        return keystoreBuilder.getKeyStore();
    }

    /**
     * Generates a {@link KeyPair} using the {@code secureRandom}
     *
     * @param secureRandom The SecureRandom
     * @return The newly generated KeyPair
     * @throws NoSuchAlgorithmException
     */
    public static KeyPair generateRSAKeyPair(final SecureRandom secureRandom)
            throws NoSuchAlgorithmException {
        final String keyType = "RSA";
        final int defaultRSAKeySize = 3072;
        final KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(keyType);
        keyPairGenerator.initialize(defaultRSAKeySize, secureRandom);
        //final String sigAlgName = "SHA384withRSA";
        final KeyPair pair = keyPairGenerator.generateKeyPair();
        final PublicKey publicKey = pair.getPublic();
        // publicKey's format must be X.509 otherwise
        if (!"X.509".equalsIgnoreCase(publicKey.getFormat())) {
            throw new IllegalArgumentException("Public key format " + publicKey.getFormat()
                    + ", isn't X.509");
        }
        return pair;
    }

    /**
     * Generates a X509 certificate using the passed {@code keyPair}
     *
     * @param keyPair         The KeyPair to use
     * @param secureRandom    The SecureRandom
     * @param subjectName     The subject to use in the certificate
     * @param subjectAltNames Any subject alternate names to use in the certificate
     * @return The newly generated certificate
     * @throws Exception
     */
    public static X509Certificate generateCert(final KeyPair keyPair, final SecureRandom secureRandom,
                                               final String subjectName, final String... subjectAltNames)
            throws Exception {
        final X500Name subject = new X500Name(subjectName, DEFAULT_SUBJECT_OU, DEFAULT_SUBJECT_ON,
                DEFAULT_SUBJECT_COUNTRY);
        final X500Name issuer = subject; // self-signed cert
        final GeneralNames generalNames;
        if (subjectAltNames == null) {
            generalNames = null;
        } else {
            generalNames = new GeneralNames();
            generalNames.add(new GeneralName(new DNSName(subjectName, true)));
            for (final String san : subjectAltNames) {
                if (san == null) {
                    continue;
                }
                final DNSName dnsName = new DNSName(san, true);
                generalNames.add(new GeneralName(dnsName));
            }
        }
        return generateCert(keyPair, secureRandom, subject, issuer, generalNames);
    }

    private static X509Certificate generateCert(final KeyPair keyPair, final SecureRandom secureRandom,
                                                final X500Name subjectName, final X500Name issuerName,
                                                final GeneralNames subjectAltNames)
            throws Exception {
        final X509CertInfo certInfo = new X509CertInfo();
        certInfo.setVersion(new CertificateVersion(CertificateVersion.V3));
        certInfo.setSerialNumber(CertificateSerialNumber.newRandom64bit(secureRandom));
        certInfo.setSubject(subjectName);
        final PublicKey publicKey = keyPair.getPublic();
        certInfo.setKey(new CertificateX509Key(publicKey));
        certInfo.setValidity(certDuration());
        certInfo.setIssuer(issuerName);
        if (subjectAltNames != null && !subjectAltNames.isEmpty()) {
            final SubjectAlternativeNameExtension sanExtn = new SubjectAlternativeNameExtension(
                    true, subjectAltNames);
            final CertificateExtensions certExtensions = new CertificateExtensions();
            certExtensions.setExtension(sanExtn.getId(), sanExtn);
            certInfo.setExtensions(certExtensions);
        }
        final PrivateKey privateKey = keyPair.getPrivate();
        final String sigAlgName = "SHA384withRSA";
        return X509CertImpl.newSigned(certInfo, privateKey, sigAlgName);
    }

    private static CertificateValidity certDuration() {
        // create a cert with 1 day validity, starting from 1 minute back
        final long currentTime = System.currentTimeMillis();
        final long oneMinuteInThePast = currentTime - (60 * 1000);
        final long oneDayInTheFuture = currentTime + (24 * 60 * 60 * 1000);
        final Date startDate = new Date(oneMinuteInThePast);
        final Date expiryDate = new Date(oneDayInTheFuture);
        return new CertificateValidity(startDate, expiryDate);
    }

    /**
     * Creates a {@link SSLContext} which is
     * {@link SSLContext#init(KeyManager[], TrustManager[], SecureRandom) initialized} using the
     * {@link javax.net.ssl.KeyManager}s and {@link javax.net.ssl.TrustManager}s available in the
     * {@code keyStore}. The {@code keyStore} is expected to be password-less.
     *
     * @param keyStore The password-less keystore
     * @return the SSLContext
     * @throws Exception
     */
    public static SSLContext createSSLContext(final KeyStore keyStore) throws Exception {
        Objects.requireNonNull(keyStore);
        final char[] password = null;
        final KeyManagerFactory kmf = KeyManagerFactory.getInstance("PKIX");
        kmf.init(keyStore, password);

        final TrustManagerFactory tmf = TrustManagerFactory.getInstance("PKIX");
        tmf.init(keyStore);

        final String protocol = "TLS";
        final SSLContext ctx = SSLContext.getInstance(protocol);
        ctx.init(kmf.getKeyManagers(), tmf.getTrustManagers(), null);
        return ctx;
    }
}
